package DataStructure.tree;


//平衡二叉查找树
public class AVLTree<K extends Comparable<K>,V> {
    private Node<K,V> root;//根节点
    private int size;//节点数量

    //节点类
    private static class Node<K extends Comparable<K>,V>{
        private K key;
        private V value;
        private Node<K,V> left;
        private Node<K,V> right;
        private int h;//以该节点为根的树的高度
        public Node(K key,V value,Node<K,V> left,Node<K,V> right) {
            this.key=key;
            this.value=value;
            this.left=left;
            this.right=right;
            this.h=1;
        }
    }

    //树的大小
    public int size(){
        return size;
    }
    //判空
    public boolean isEmpty(){
        return size==0;
    }
    //树的高度
    public int height(Node<K,V> n){
        return n==null?0:n.h;
    }
    //添加节点
    public void put(K key,V value){
        root=put(root,key,value);
    }
    private Node<K,V> put(Node<K,V> n,K key,V value){
        if(n==null) {
            size++;//节点为空
            return new Node<>(key, value, null, null);
        }
        int cmp=key.compareTo(n.key);
        if(cmp>0){
            n.right=put(n.right,key,value);
        } else if(cmp<0){
            n.left=put(n.left,key,value);
        }else
            n.value=value;//存在key,直接修改
        return rotate(n);
    }
    //查找节点
    public V get(K key){
        return get(root,key);
    }
    private V get(Node<K,V> n,K key){
        if(n==null)//不存在
            return null;
        int cmp=key.compareTo(n.key);
        if(cmp>0)
            return get(n.right,key);
        else if(cmp<0)
            return get(n.left,key);
        else
            return n.value;//已存在
    }
    //是否包含节点
    public boolean contains(K key){
        return get(key)!=null;
    }
    //删除节点
    public Node<K,V> remove(K key){
        return remove(root,key);
    }
    private Node<K,V> remove(Node<K,V> n,K key){
        if(n==null)//不存在
            return null;
        //查找key对应节点
        int cmp=key.compareTo(n.key);
        if(cmp>0){
            n.right=remove(n.right,key);
        } else if(cmp<0){
            n.left=remove(n.left,key);
        } else{
            size--;
            if(n.left==null&&n.right==null)
                return null;
            else if(n.right==null)
                n=n.left;
            else if(n.left==null)
                n=n.right;
            else{
                //查找右子树最小节点
                Node<K,V> min=min(n.right);
                //替换当前节点
                n.key=min.key;
                n.value=min.value;
                n.right=remove(n.right,min.key);//删除min
            }
        }
        return rotate(n);
    }
    private Node<K,V> rotate(Node<K,V> n) {
        //左子树较高
        if(height(n.left)-height(n.right)>=2){
            if(height(n.left.left)>height(n.left.right))
                return rightRotate(n);//LL
            return leftRightRotate(n);//LR
        }else if(height(n.right)-height(n.left)>=2){//右子树较高
            if(height(n.right.right)>height(n.right.left))
                return leftRotate(n);//RR
            return rightLeftRotate(n);//RL
        }
        n.h=Math.max(height(n.left),height(n.right))+1;//更新高度
        return n;
    }
    private Node<K,V> leftRightRotate(Node<K,V> n) {
        n.left=leftRotate(n.left);
        return rightRotate(n);
    }
    private Node<K,V> rightLeftRotate(Node<K,V> n) {
        n.right=rightRotate(n.right);
        return leftRotate(n);
    }
    private Node<K,V> rightRotate(Node<K,V> n) {
        Node<K,V> ans=n.left;
        n.left=ans.right;
        ans.right=n;
        n.h=Math.max(height(n.left),height(n.right))+1;
        ans.h=Math.max(height(ans.left),n.h)+1;
        return ans;
    }
    private Node<K,V> leftRotate(Node<K,V> n) {
        Node<K,V> ans=n.right;
        n.right=ans.left;
        ans.left=n;
        n.h=Math.max(height(n.left),height(n.right))+1;
        ans.h=Math.max(height(ans.right),n.h)+1;
        return ans;
    }
    //返回最小键
    public K min(){
        return min(root).key;
    }
    private Node<K,V> min(Node<K,V> n){
        if(n.left!=null)
            return min(n.left);
        return n;
    }
    //返回最大键
    public K max(){
        return max(root).key;
    }
    private Node<K,V> max(Node<K,V> n){
        if(n.right!=null)
            return max(n.right);
        return n;
    }
}
